package org.xyl.example;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.xyl.engine.ChromaInitializer;

import java.util.List;

public class ChromaExample {

    public static void main(String[] args) {
        // 1. 初始化 Chroma EmbeddingStore
        EmbeddingStore<TextSegment> embeddingStore = ChromaInitializer.initializeChroma();


        // 2. 初始化嵌入模型 (这里使用内置的小型模型)
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        // 3. 创建嵌入向量并存储
        TextSegment segment1 = TextSegment.from("LangChain4J 是一个 Java 的 LLM 集成框架");
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        embeddingStore.add(embedding1, segment1);

        TextSegment segment2 = TextSegment.from("Chroma 是一个轻量级的向量数据库");
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        embeddingStore.add(embedding2, segment2);

        // 4. 相似性搜索
        String query = "Java 的 LLM 框架";
        Embedding queryEmbedding = embeddingModel.embed(query).content();

        List<EmbeddingMatch<TextSegment>> matches = embeddingStore.findRelevant(queryEmbedding, 1);

        // 5. 输出结果
        matches.forEach(match -> {
            System.out.printf("相似度: %.2f%%\n", match.score() * 100);
            System.out.println("匹配文本: " + match.embedded().text());
        });
    }
}